from service.rag_neo4j_service import RagNeo4jService
import json


class CompanyNodeLLMService:
    def __init__(self):
        self.ragNeo4jService = RagNeo4jService()
        self.chain = self.ragNeo4jService.create_chain()

    # 处理json格式的字符串
    def handle_json(self, json_str):
        if "```json" in json_str:
            text = json_str.split("```json")[1].split("```")[0]
        else:
            text = json_str
        # 转换成json格式
        info = json.loads(text)
        return info

    # 创建节点
    def create_node(self, **kwargs):
        info = f"创建一个企业节点,它的属性有：{kwargs}"
        self.chain.invoke({"query": info})

    # 查询节点
    def query_node(self, **kwargs):
        if len(kwargs) > 0:
            info = f"查询一个企业节点,它的属性有：{kwargs},返回的数据必须是json格式"
        else:
            info = "查询所有企业节点,返回的数据必须是json格式"
        rs = self.chain.invoke({"query": info})
        data = self.handle_json(rs["result"])
        print(data)
        return data

    # 删除节点
    def delete_node(self, **kwargs):
        info = f"删除一个企业节点,它的属性有：{kwargs}"
        self.chain.invoke({"query": info})

    # 更新节点
    def update_node(self, **kwargs):
        info = f"更新一个企业节点,它的属性有：{kwargs}"
        self.chain.invoke({"query": info})


if __name__ == '__main__':
    company_node_service = CompanyNodeLLMService()
    # new_company_attrs = {
    #     "上市代码": "002594",
    #     "企业名称": "比亚迪股份有限公司",
    #     "信用评级": "AAA（大公国际）",
    #     "报告期": "2025-06-30",
    #     "主营业务": "新能源汽车研发、制造与销售；动力电池生产；光伏产品研发",
    #     "成立日期": "1995-02-10",
    #     "注册资本": 731142.44,
    #     "上市状态": "是（A+H股上市）",
    #     "统一社会信用代码": "91440300192317458F",
    #     "市盈率": 25.8
    # }
    #
    # company_node_service.create_node(**new_company_attrs)
    # print("新企业节点创建请求已发送")
    # company_node_service.query_node()
